# src/evaluate.py

import torch
from torch.utils.data import DataLoader
from data_loader import ECGDataset
from model_wavelet import ECG_Autoencoder_Classifier_Lstm, ECG_Autoencoder_Classifier_Transformer, ECG_Autoencoder_Classifier_Mamba, ECG_Autoencoder_Classifier_Cnn
from model_time import ChannelWiseCNN
from sklearn.metrics import classification_report, roc_auc_score, confusion_matrix
import pandas as pd
import os
from datetime import datetime
import numpy as np
import json
from collections import defaultdict
import seaborn as sns
import matplotlib.pyplot as plt

def load_all_models(base_path, config, device):
    """加载所有预训练模型"""
    models = {}
    
    model_classes = {
        'WaveletLstm': ECG_Autoencoder_Classifier_Lstm,
        'WaveletTransformer': ECG_Autoencoder_Classifier_Transformer,
        #'WaveletMamba': ECG_Autoencoder_Classifier_Mamba,
        #'WaveletCnn': ECG_Autoencoder_Classifier_Cnn,
        'TimeCNN': ChannelWiseCNN
    }
    
    for model_name in model_classes.keys():
        model_path = os.path.join(base_path, model_name, 'model_weights.pth')
        if not os.path.exists(model_path):
            print(f"找不到模型文件: {model_path}")
            continue
            
        model = model_classes[model_name](
            input_dim=config['input_dim'],
            hidden_dim=config['hidden_dim'],
            latent_dim=config['latent_dim'],
            num_layers=config['num_layers']
        )
        model.load_state_dict(torch.load(model_path, map_location=device))
        model = model.to(device)
        model.eval()
        models[model_name] = model
    
    return models

def plot_confusion_matrix(cm, labels, title, save_path):
    """绘制并保存混淆矩阵热力图"""
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=labels, yticklabels=labels)
    plt.title(title)
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()

def evaluate_single_model(model, dataloader, device, model_name, result_dir):
    """评估单个模型的性能"""
    preds_list = []
    labels_list = []
    probs_list = []
    sample_ids = []

    with torch.no_grad():
        for batch in dataloader:
            inputs, labels, ids = batch
            inputs = inputs.to(device)
            labels = labels.to(device).float()

            logits = model(inputs)
            probs = torch.sigmoid(logits)
            preds = (probs >= 0.55).float()

            preds_list.extend(preds.cpu().numpy())
            probs_list.extend(probs.cpu().numpy())
            labels_list.extend(labels.cpu().numpy())
            sample_ids.extend(ids)

    # 展平结果
    preds_list = np.array(preds_list).flatten()
    probs_list = np.array(probs_list).flatten()
    labels_list = np.array(labels_list).flatten()

    # 计算准确率
    accuracy = (preds_list == labels_list).mean()

    # 生成评估报告
    report = classification_report(labels_list, preds_list, output_dict=True)
    report_text = classification_report(labels_list, preds_list)
    auc = roc_auc_score(labels_list, preds_list)

    # 保存结果
    model_dir = os.path.join(result_dir, model_name)
    os.makedirs(model_dir, exist_ok=True)

    # 保存评估报告
    with open(os.path.join(model_dir, 'evaluation_report.txt'), 'w') as f:
        f.write(report_text)
        f.write(f'\nROC-AUC Score: {auc:.4f}')
        f.write(f'\nAccuracy: {accuracy:.4f}')

    # 保存详细预测结果
    results_df = pd.DataFrame({
        'Sample_ID': sample_ids,
        'True_Labels': labels_list,
        'Predictions': preds_list,
        'Probabilities': probs_list
    })
    results_df['sort_key'] = results_df['Sample_ID'].str.extract('(\d+)').astype(int)
    results_df = results_df.sort_values('sort_key').drop('sort_key', axis=1)
    results_df.to_csv(os.path.join(model_dir, 'predictions.csv'), index=False)

    # 计算混淆矩阵
    cm = confusion_matrix(labels_list, preds_list)
    labels = ['Negative', 'Positive']
    
    # 绘制并保存混淆矩阵
    plot_confusion_matrix(
        cm, 
        labels, 
        f'Confusion Matrix - {model_name}',
        os.path.join(model_dir, 'confusion_matrix.png')
    )

    # 将混淆矩阵数据保存到文本文件
    with open(os.path.join(model_dir, 'confusion_matrix.txt'), 'w') as f:
        f.write(f"Confusion Matrix for {model_name}:\n")
        f.write("True Negative: {}\n".format(cm[0][0]))
        f.write("False Positive: {}\n".format(cm[0][1]))
        f.write("False Negative: {}\n".format(cm[1][0]))
        f.write("True Positive: {}\n".format(cm[1][1]))

    return probs_list, labels_list, auc, accuracy

def ensemble_evaluate(models, dataloader, device):
    """集成评估所有模型"""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    result_dir = os.path.join('results/ensemble_evaluation', timestamp)
    os.makedirs(result_dir, exist_ok=True)

    # 存储每个模型的准确率和预测概率
    model_accuracies = {}
    model_probs = {}
    true_labels = None

    # 评估每个模型
    for model_name, model in models.items():
        print(f"\n评估 {model_name} 模型...")
        probs, labels, auc, accuracy = evaluate_single_model(model, dataloader, device, model_name, result_dir)
        model_accuracies[model_name] = accuracy
        model_probs[model_name] = probs
        if true_labels is None:
            true_labels = labels

    # 计算基于准确率的权重
    total_accuracy = sum(model_accuracies.values())
    weights = {name: acc/total_accuracy for name, acc in model_accuracies.items()}

    # 计算加权集成预测
    ensemble_probs = np.zeros_like(true_labels, dtype=float)
    for model_name in models.keys():
        ensemble_probs += weights[model_name] * model_probs[model_name]

    # 生成集成预测结果
    ensemble_preds = (ensemble_probs >= 0.55).astype(float)
    
    # 评估集成结果
    ensemble_report = classification_report(true_labels, ensemble_preds)
    ensemble_auc = roc_auc_score(true_labels, ensemble_preds)
    ensemble_accuracy = (ensemble_preds == true_labels).mean()

    # 保存集成结果
    with open(os.path.join(result_dir, 'ensemble_report.txt'), 'w') as f:
        f.write("集成模型评估报告:\n")
        f.write("模型权重 (基于准确率):\n")
        for model_name, weight in weights.items():
            f.write(f"{model_name}: {weight:.4f}\n")
        f.write("\n分类报告:\n")
        f.write(ensemble_report)
        f.write(f'\nROC-AUC Score: {ensemble_auc:.4f}')
        f.write(f'\nAccuracy: {ensemble_accuracy:.4f}')

    # 计算集成模型的混淆矩阵
    ensemble_cm = confusion_matrix(true_labels, ensemble_preds)
    labels = ['Negative', 'Positive']
    
    # 绘制并保存集成模型的混淆矩阵
    plot_confusion_matrix(
        ensemble_cm, 
        labels, 
        'Confusion Matrix - Ensemble Model',
        os.path.join(result_dir, 'ensemble_confusion_matrix.png')
    )

    # 将集成模型的混淆矩阵数据保存到文本文件
    with open(os.path.join(result_dir, 'ensemble_confusion_matrix.txt'), 'w') as f:
        f.write("Confusion Matrix for Ensemble Model:\n")
        f.write("True Negative: {}\n".format(ensemble_cm[0][0]))
        f.write("False Positive: {}\n".format(ensemble_cm[0][1]))
        f.write("False Negative: {}\n".format(ensemble_cm[1][0]))
        f.write("True Positive: {}\n".format(ensemble_cm[1][1]))

    print("\n集成模型评估报告:")
    print("模型权重 (基于准确率):", weights)
    print("\n分类报告:")
    print(ensemble_report)
    print(f'集成模型 ROC-AUC Score: {ensemble_auc:.4f}')
    print(f'集成模型 Accuracy: {ensemble_accuracy:.4f}')

def main():
    # 配置参数
    test_data_dir = 'data/processed/test_data'
    models_base_path = 'results/final_model'
    batch_size = 32
    
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 模型配置
    model_config = {
        'input_dim': 1,
        'hidden_dim': 256,
        'latent_dim': 128,
        'num_layers': 8
    }
    
    # 加载数据
    test_dataset = ECGDataset(test_data_dir)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    # 加载所有模型
    models = load_all_models(models_base_path, model_config, device)
    
    # 进行集成评估
    ensemble_evaluate(models, test_loader, device)

if __name__ == '__main__':
    main()

